
#include "tril_tiling.h"
#include "register/op_def_registry.h"


namespace optiling {
const uint32_t BLOCK_SIZE = 32;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  TrilTilingData tiling;
  uint32_t sizeofdatatype;
        uint32_t totalLengthAligned;
        uint32_t ALIGN_NUM;
        uint32_t tileLength;
        uint32_t lasttileLength;
        uint32_t tile_num;
       uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize();
        auto *pvalue = context->GetAttrs()->GetInt(0);
        auto Shape = context->GetInputTensor(0)->GetOriginShape();
        uint32_t xShape = Shape.GetDim(Shape.GetDimNum() - 2);
        uint32_t yShape = Shape.GetDim(Shape.GetDimNum() - 1);
	auto dt = context->GetInputTensor(0)->GetDataType();

        if (dt == 1)
        { // float16
            sizeofdatatype = 2;
        }
        else if (dt == 0 || dt == 3)
        { // float
            sizeofdatatype = 4;
        }

        ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;
        
 	if (totalLength % ALIGN_NUM != 0)
        { // 不对齐，先32位对齐
            totalLengthAligned =
                ((totalLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
        }
        else
        {
            totalLengthAligned = totalLength;
        }
	uint32_t ub_block_num = 4;
        context->SetBlockDim(1);
        auto block_dim = context->GetBlockDim();
     	tileLength = ALIGN_NUM* ub_block_num;
	 if ((totalLengthAligned/ ALIGN_NUM) % ub_block_num == 0)
        {
            tileLength = ALIGN_NUM* ub_block_num;
        if (tile_num == 0)
            {
                tile_num = 1;
                lasttileLength = totalLength;
            }else{

                lasttileLength = tileLength;
            }
        }
        else
        {
            tile_num = tile_num + 1;
            tileLength = ub_block_num * ALIGN_NUM;
            lasttileLength = totalLength - (tile_num - 1) * tileLength;
        }

	tiling.set_value(*pvalue);
        tiling.set_xShape(xShape);
        tiling.set_yShape(yShape);
        context->SetBlockDim(1);
	printf("blockLength: %d\n",totalLength);
	printf("tileLength: %d \n", tileLength);
        printf("lasttileLength: %d \n", lasttileLength);
        printf("tileNum: %d \n",tile_num);
        tiling.set_blockLength(totalLength);
        tiling.set_tileNum(tile_num);
        tiling.set_tileLength(tileLength);
        tiling.set_lasttileLength(lasttileLength);
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                            context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        currentWorkspace[0] = 0;
        return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Tril : public OpDef {
public:
    explicit Tril(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("diagonal").Int();

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(Tril);
}
